import matplotlib.pyplot as plt
import numpy as np
from test_4_00 import create_polar_system

# 创建极坐标系
fig, ax = create_polar_system(r_range=(0, 10))

ax.set_title('阿基米德螺线 (r = aθ)', fontsize=16, pad=20)

# 定义阿基米德螺线参数
a = 0.5  # 螺线参数
theta = np.linspace(0, 4*np.pi, 1000)  # 角度范围 0 到 4π
r = a * theta  # 阿基米德螺线方程 r = aθ

# 绘制阿基米德螺线
ax.plot(theta, r, 'b-', linewidth=2, label=f'r = {a}θ')

# 添加图例
ax.legend(loc='upper right', bbox_to_anchor=(1.2, 1.1))

# 标记一些特殊点
# 每隔π标记一个点
special_thetas = np.array([0, np.pi, 2*np.pi, 3*np.pi, 4*np.pi])
special_rs = a * special_thetas
ax.plot(special_thetas, special_rs, 'ro', markersize=6)

# 添加特殊点标签
labels = ['0', 'π', '2π', '3π', '4π']
for i, (t, r_val) in enumerate(zip(special_thetas, special_rs)):
    if t != 0:  # 不标记原点
        ax.text(t + 0.1, r_val + 0.2, labels[i], fontsize=10, color='red')

plt.tight_layout()
plt.show()